use std::{ops::Range, time::Instant};

use godot::{classes::RandomNumberGenerator, prelude::*};
use rand::{Rng, SeedableRng, rngs::SmallRng};

use crate::zmath::ZMath;

#[derive(GodotClass)]
#[class(init, base=Object)]
pub struct ZUtil {}

#[godot_api]
impl ZUtil {
    #[func]
    pub fn for_each_range(a: i64, b: i64, step: i64, body: Callable) {
        let (mut min, max) = ZMath::minmax_raw(a, b);
        while min < max {
            body.call(&[Variant::from(min)]);
            min += step
        }
    }

    #[func]
    pub fn for_each(collection: Variant, body: Callable) {
        match collection.get_type() {
            VariantType::ARRAY => match body.get_argument_count() {
                1 => {
                    for i in collection.to::<VariantArray>().iter_shared() {
                        body.call(&[i]);
                    }
                }
                2 => {
                    for (idx, v) in collection.to::<VariantArray>().iter_shared().enumerate() {
                        body.call(&[Variant::from(idx as i64), v]);
                    }
                }
                _ => godot_error!("Too many args!"),
            },
            VariantType::DICTIONARY => {
                if body.get_argument_count() != 2 {
                    godot_error!("Expect 2 args.")
                }

                for (k, v) in collection.to::<Dictionary>().iter_shared() {
                    body.call(&[k, v]);
                }
            }
            _ => godot_error!("Unable to traverse a non-collection type"),
        }
    }

    #[func]
    pub fn retain_arr(mut arr: VariantArray, filter: Callable) {
        let mut del = 0;
        for i in 0..arr.len() {
            let cur = arr.at(i);
            if !filter.call(&[cur.clone()]).booleanize() {
                del += 1
            } else if del > 0 {
                let t = arr.at(i - del);
                arr.set(i - del, &cur);
                arr.set(i, &t);
            }
        }
        arr.resize(arr.len() - del, &Variant::nil());
    }

    #[func]
    pub fn is_mask(flags: i64, mask: i64) -> bool {
        flags & mask == mask
    }

    #[func]
    pub fn except_mask(flags: i64, mask: i64) -> bool {
        flags & mask == flags
    }

    #[func]
    pub fn any_flag(flags: i64, mask: i64) -> bool {
        flags & mask != 0
    }

    #[func]
    pub fn null_or(val: Variant, default: Variant) -> Variant {
        if val.is_nil() { default } else { val }
    }

    #[func]
    pub fn and_then(val: Variant, then: Callable) -> Variant {
        if val.is_nil() { val } else { then.call(&[val]) }
    }

    #[func]
    pub fn car(val: Variant) -> Variant {
        match val.get_type() {
            VariantType::ARRAY => val.to::<VariantArray>().pop_front().unwrap_or_default(),
            _ => val,
        }
    }

    #[func]
    pub fn cdr(val: Variant) -> Variant {
        match val.get_type() {
            VariantType::ARRAY => {
                val.to::<VariantArray>().pop_front();
                val
            }
            _ => Variant::nil(),
        }
    }

    #[func]
    pub fn cadr(val: Variant) -> Variant {
        Self::car(Self::cdr(val))
    }

    fn rand_in_weights_raw(
        mut rand_range: impl FnMut(Range<i64>) -> i64,
        weights: impl Iterator<Item = i64>,
    ) -> Option<usize> {
        let mut candidate = vec![];
        let mut total = 0;

        for w in weights {
            total += w;
            candidate.push(total);
        }

        let pick = rand_range(0..total);

        candidate
            .into_iter()
            .enumerate()
            .find(|(_, t)| *t > pick)
            .map(|x| x.0)
    }

    #[func]
    pub fn rand_in_weights(weights: VariantArray) -> i64 {
        match weights.len() {
            0 => return -1,
            1 => return 0,
            _ => {}
        }
        
        let mut rng = SmallRng::from_os_rng();
        Self::rand_in_weights_raw(
            |range| rng.random_range(range),
            weights.iter_shared().map(|w| w.to::<i64>()),
        )
        .map(|x| x as i64)
        .unwrap_or(-1)
    }

    #[func]
    pub fn rand_in_weights_rng(mut rng: Gd<RandomNumberGenerator>, weights: VariantArray) -> i64 {
        match weights.len() {
            0 => return -1,
            1 => return 0,
            _ => {}
        }

        Self::rand_in_weights_raw(
            |range| rng.randi_range(range.start as _, range.end as _) as _,
            weights.iter_shared().map(|w| w.to::<i64>()),
        )
        .map(|x| x as i64)
        .unwrap_or(-1)
    }

    #[func]
    pub fn bench(block: Callable, times: i64) -> f64 {
        let begin = Instant::now();
        for _ in 0..times {
            block.call(&[]);
        }
        (Instant::now() - begin).as_secs_f64() * 1000.0 / times as f64
    }
}

#[cfg(test)]
mod test {
    use std::time::Instant;

    use rand::{Rng, SeedableRng, rngs::SmallRng};

    use crate::zutil::ZUtil;

    #[test]
    fn test_rand_w() {
        let mut rng = SmallRng::from_os_rng();
        let pick_times = 10;
        let rand_times = 100000;

        let instant = Instant::now();
        for _ in 0..pick_times {
            let mut rate = [0.0, 0.0, 0.0];
            let weights = [
                rng.random_range(1..100),
                rng.random_range(1..100),
                rng.random_range(1..100),
            ];
            let total = weights.iter().cloned().sum::<i64>();
            for _ in 0..rand_times {
                let i = ZUtil::rand_in_weights_raw(
                    |range| rng.random_range(range),
                    weights.iter().cloned(),
                );

                rate[i.unwrap()] += 1.0;
            }

            rate.iter_mut().for_each(|x| *x /= rand_times as f64);

            let mut standard = [0.0; 3];
            for i in 0..weights.len() {
                standard[i] = weights[i] as f64 / total as f64;
                assert!((standard[i] - rate[i]).abs() < 0.01)
            }

            println!("{weights:?} {rate:?} {standard:?}")
        }

        println!(
            "average: {}ms",
            (Instant::now() - instant).as_secs_f64() * 1000.0 / ((pick_times * rand_times) as f64)
        )
    }
}
